# -*- coding:utf-8 -*-
import numpy as np
import time
import os

from sko.GA import GA_TSP # 遗传算法库

from preprocess_data import get_data

features,labels = get_data('训练集',shuffle=False)
# features : [num_sample,num_feature]

# 每类图片选取一个作为代表
features = features[::144,:]
num_feature = features.shape[1]
# width : 图片的宽,从440的所有公因数(1,2,4,5,8,10,11,20,22,40,44,55,88,110,220,440)选取
width = 440
assert num_feature % width == 0
height = num_feature//width

print(features.shape)
if not os.path.exists('saved_models/d_features.npz'):
    d_features = np.zeros((200,num_feature,num_feature),dtype='int32')
    for i in range(200):
        for j in range(num_feature):
            for k in range(num_feature):
                d_features[i,j,k] = np.abs(features[i,j] - features[i,k])

    # 下面的这个以后可以用max来试一下
    d_features = d_features.mean(axis = 0)

    np.savez('saved_models/d_features.npz',d_features = d_features)
else:
    d_features = np.load('saved_models/d_features.npz')['d_features']

def cost(i,routine):
    cost_i = 0
    if i - width >= 0:
        cost_i += d_features[routine[i-width],routine[i]]
    if i + width < routine.shape[0]:
        cost_i += d_features[routine[i+width],routine[i]]
    if i % width != 0 and i-1>=0:
        cost_i += d_features[routine[i-1],routine[i]]
    if i % width != width-1 and i+1<routine.shape[0]:
        cost_i += d_features[routine[i+1],routine[i]]
    return cost_i

# 下面把这个问题抽象成TSP
def cal_total_cost(routine):
    '''The objective function. input routine, return total distance.
    cal_total_cost(np.arange(num_points))
    '''
    # routine[i] : 最优路径第i个点的下标
    return np.sum([cost(i,routine) for i in range(routine.shape[0])])

routine = np.array([i for i in range(num_feature)])
print(cal_total_cost(routine)/((height-1)*width+(width-1)*height))

t_start = time.time()
ga_tsp = GA_TSP(func=cal_total_cost, n_dim=num_feature, size_pop=50, max_iter=500, prob_mut=1)
best_points, best_distance = ga_tsp.run()

# 每个宽度都保存一个模型
np.savez('saved_models/select/width-%s.npz'%width,best_points = np.array(best_points))

print(best_points)
print(best_distance/((height-1)*width+(width-1)*height))
t_end = time.time()
print('time:%.2f'%(t_end-t_start))